import os
import requests
import json
import time
from dotenv import load_dotenv
from PIL import Image  # Add Pillow for image processing

# Load environment variables from .env file
load_dotenv()

# Get API key from environment variable
api_key = os.getenv("DASHSCOPE_API_KEY")

# Define the DashScope API endpoint
DASHSCOPE_API_URL = "https://dashscope.aliyuncs.com/api/v1/services/aigc/text2image/image-synthesis"

# Define the target directories
TARGET_DIR = "static/shikong"
STATIC_DIR = "static"
TABBAR_DIR = "static/tabbar"

# Create the directories if they don't exist
os.makedirs(TARGET_DIR, exist_ok=True)
os.makedirs(STATIC_DIR, exist_ok=True)
os.makedirs(TABBAR_DIR, exist_ok=True)

# Define the image specifications for each image
image_specs = [
    {
        "filename": "guanzhong.jpg",
        "prompt": "古代中国春秋时期的政治家管仲的肖像画，他有威严的神情、传统中国式胡须，穿着春秋时期的正式官服。风格应为写实的历史人物肖像，具有中国传统绘画元素。",
        "description": "管仲，春秋时期齐国宰相，莒县人",
        "directory": TARGET_DIR
    },
    {
        "filename": "liuxie.jpg",
        "prompt": "南北朝时期文学批评家刘勰的肖像画，他看起来学识渊博且睿智，穿着南北朝时期的传统服饰。肖像应体现其文学家的气质，周围有《文心雕龙》卷轴或书籍等象征其文学理论成就的元素。",
        "description": "刘勰，南北朝文学批评家，《文心雕龙》作者",
        "directory": TARGET_DIR
    },
    {
        "filename": "zuosi.jpg",
        "prompt": "西晋文学家诗人左思的肖像画，他有一种优雅、艺术的气质，身着晋代传统士人服装。肖像应捕捉其诗人气质，采用优雅的笔触，具有中国传统绘画风格。",
        "description": "左思，西晋诗人文学家，莒县人",
        "directory": TARGET_DIR
    },
    {
        "filename": "jiangliang.jpg",
        "prompt": "明代诗人书法家姜亮的肖像画，他有一种内敛而高雅的气质，身着明朝士人服饰。肖像应展现他与书法毛笔或卷轴在一起的场景，采用适合明代时期的中国传统绘画风格。",
        "description": "姜亮，明代诗人书法家，莒县人",
        "directory": TARGET_DIR
    },
    {
        "filename": "user-avatar.png",
        "prompt": "简约现代风格的用户头像图标，使用蓝色调，设计成圆形，内部有一个简洁的人形轮廓。图标应当干净、专业，适合在应用程序界面中作为默认用户头像使用。",
        "description": "用户默认头像",
        "directory": STATIC_DIR
    },
    {
        "filename": "default-avatar.png",
        "prompt": "一个简约、中性的默认头像设计，灰色调，圆形，包含一个通用的人形轮廓。设计应当简洁、专业，没有性别特征，适合作为网站或应用的默认头像。",
        "description": "通用默认头像",
        "directory": STATIC_DIR
    },
    {
        "filename": "assistant-avatar.png",
        "prompt": "AI助手头像，采用高科技风格设计，紫色调，圆形头像中心有一个抽象的机器人或AI符号。设计要现代、友好、智能感强，适合作为智能助手的代表形象。",
        "description": "AI助手头像",
        "directory": STATIC_DIR
    },
    # Tabbar icons - normal state
    {
        "filename": "chat.png",
        "prompt": "设计一个简约的聊天图标，灰色调，线条风格，适合移动应用的标签栏。图标应当简洁清晰，大小为64x64像素，背景透明，适合深色和浅色背景使用。",
        "description": "智问莒县标签图标",
        "directory": TABBAR_DIR
    },
    {
        "filename": "clock.png",
        "prompt": "设计一个简约的时钟图标，灰色调，线条风格，适合移动应用的标签栏。图标应当简洁清晰，大小为64x64像素，背景透明，适合深色和浅色背景使用。",
        "description": "时空对话标签图标",
        "directory": TABBAR_DIR
    },
    {
        "filename": "image.png",
        "prompt": "设计一个简约的图片图标，灰色调，线条风格，适合移动应用的标签栏。图标应当简洁清晰，大小为64x64像素，背景透明，适合深色和浅色背景使用。",
        "description": "传图讲解标签图标",
        "directory": TABBAR_DIR
    },
    {
        "filename": "location.png",
        "prompt": "设计一个简约的位置标记图标，灰色调，线条风格，适合移动应用的标签栏。图标应当简洁清晰，大小为64x64像素，背景透明，适合深色和浅色背景使用。",
        "description": "打卡有礼标签图标",
        "directory": TABBAR_DIR
    },
    {
        "filename": "note.png",
        "prompt": "设计一个简约的笔记本图标，灰色调，线条风格，适合移动应用的标签栏。图标应当简洁清晰，大小为64x64像素，背景透明，适合深色和浅色背景使用。",
        "description": "游记生成标签图标",
        "directory": TABBAR_DIR
    },
    # Tabbar icons - active state
    {
        "filename": "chat-active.png",
        "prompt": "设计一个简约的聊天图标，蓝色调（#007aff），线条风格，适合移动应用的标签栏选中状态。图标应当简洁清晰，大小为64x64像素，背景透明，适合深色和浅色背景使用。",
        "description": "智问莒县标签选中图标",
        "directory": TABBAR_DIR
    },
    {
        "filename": "clock-active.png",
        "prompt": "设计一个简约的时钟图标，蓝色调（#007aff），线条风格，适合移动应用的标签栏选中状态。图标应当简洁清晰，大小为64x64像素，背景透明，适合深色和浅色背景使用。",
        "description": "时空对话标签选中图标",
        "directory": TABBAR_DIR
    },
    {
        "filename": "image-active.png",
        "prompt": "设计一个简约的图片图标，蓝色调（#007aff），线条风格，适合移动应用的标签栏选中状态。图标应当简洁清晰，大小为64x64像素，背景透明，适合深色和浅色背景使用。",
        "description": "传图讲解标签选中图标",
        "directory": TABBAR_DIR
    },
    {
        "filename": "location-active.png",
        "prompt": "设计一个简约的位置标记图标，蓝色调（#007aff），线条风格，适合移动应用的标签栏选中状态。图标应当简洁清晰，大小为64x64像素，背景透明，适合深色和浅色背景使用。",
        "description": "打卡有礼标签选中图标",
        "directory": TABBAR_DIR
    },
    {
        "filename": "note-active.png",
        "prompt": "设计一个简约的笔记本图标，蓝色调（#007aff），线条风格，适合移动应用的标签栏选中状态。图标应当简洁清晰，大小为64x64像素，背景透明，适合深色和浅色背景使用。",
        "description": "游记生成标签选中图标",
        "directory": TABBAR_DIR
    }
]

def file_needs_generation(directory, filename):
    """Check if a file needs to be generated (doesn't exist or has zero size)"""
    # Handle filenames with or without leading slash
    if filename.startswith('/'):
        filename = filename[1:]
    
    file_path = os.path.join(directory, filename)
    # Check if file exists and has size greater than 0
    return not (os.path.exists(file_path) and os.path.getsize(file_path) > 0)

def generate_image(prompt, directory, filename):
    """Generate an image using DashScope API and save it to the specified path."""
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}",
        "X-DashScope-Async": "enable"
    }
    
    data = {
        "model": "wanx2.1-t2i-turbo",
        "input": {
            "prompt": prompt
        },
        "parameters": {
            "size": "512*512",
            "n": 1
        }
    }
    
    try:
        print(f"Generating image for: {filename}")
        response = requests.post(DASHSCOPE_API_URL, headers=headers, json=data)
        
        if response.status_code == 200:
            result = response.json()
            task_id = result.get('output', {}).get('task_id')
            
            if not task_id:
                print("Error: No task_id found in response")
                print(f"Response: {result}")
                return False
            
            # Check the task status
            task_status_url = f"https://dashscope.aliyuncs.com/api/v1/tasks/{task_id}"
            
            # Poll for task completion
            for _ in range(30):  # Try for up to 30 times (5 minutes)
                time.sleep(10)  # Wait 10 seconds between checks
                
                status_response = requests.get(
                    task_status_url, 
                    headers={"Authorization": f"Bearer {api_key}"}
                )
                
                if status_response.status_code == 200:
                    status_data = status_response.json()
                    task_status = status_data.get('output', {}).get('task_status')
                    
                    if task_status == "SUCCEEDED":
                        # Get the image URL
                        image_urls = status_data.get('output', {}).get('results', [])
                        
                        if image_urls and 'url' in image_urls[0]:
                            image_url = image_urls[0]['url']
                            
                            # Download the image
                            image_response = requests.get(image_url)
                            if image_response.status_code == 200:
                                file_path = os.path.join(directory, filename)
                                with open(file_path, "wb") as file:
                                    file.write(image_response.content)
                                print(f"Successfully saved image to {file_path}")
                                return True
                            else:
                                print(f"Failed to download image: {image_response.status_code}")
                        else:
                            print("No image URL found in response")
                            print(f"Response: {status_data}")
                            return False
                    
                    elif task_status == "FAILED":
                        print(f"Task failed: {status_data.get('output', {}).get('message', 'Unknown error')}")
                        return False
                    
                    # Continue polling if still pending or running
                    print(f"Task status: {task_status}, waiting...")
                    
                else:
                    print(f"Failed to check task status: {status_response.status_code}")
                    print(f"Response: {status_response.text}")
                    return False
            
            print("Timeout waiting for task to complete")
            return False
        else:
            print(f"Failed to generate image: {response.status_code}")
            print(f"Response: {response.text}")
            return False
    except Exception as e:
        print(f"Error generating image: {str(e)}")
        return False

def compress_image(file_path, target_size=(128, 128)):
    """Compress and resize an image to the target size"""
    try:
        # Open the image
        with Image.open(file_path) as img:
            # If image has transparency (PNG), preserve it
            if img.mode in ('RGBA', 'LA') or (img.mode == 'P' and 'transparency' in img.info):
                # Convert to RGBA if it's not already
                img = img.convert('RGBA')
                
                # Create a blank image with transparent background
                new_img = Image.new('RGBA', target_size, (0, 0, 0, 0))
                
                # Resize the original image with high quality and anti-aliasing
                resized_img = img.resize(target_size, Image.LANCZOS)
                
                # Paste the resized image onto the new transparent image
                new_img.paste(resized_img, (0, 0), resized_img)
                
                # Save the image with transparency
                new_img.save(file_path, optimize=True)
            else:
                # For images without transparency
                img = img.resize(target_size, Image.LANCZOS)
                img.save(file_path, optimize=True)
                
            print(f"Successfully compressed {file_path} to {target_size[0]}x{target_size[1]}")
            return True
    except Exception as e:
        print(f"Error compressing image {file_path}: {str(e)}")
        return False

def get_target_size(directory):
    """Determine the target size based on the image directory"""
    if directory == TABBAR_DIR:
        return (64, 64)  # Tabbar icons are 128x128
    else:
        return (256, 256)  # Other images (avatars, historical figures) are 256x256

def main():
    print("Starting image generation for historical figures, avatars and tabbar icons...")
    
    if not api_key:
        print("Error: DASHSCOPE_API_KEY not set in environment variables or .env file")
        print("Please set your DashScope API key before running this script")
        return
    
    # Create the target directories if they don't exist
    os.makedirs(TARGET_DIR, exist_ok=True)
    os.makedirs(STATIC_DIR, exist_ok=True)
    os.makedirs(TABBAR_DIR, exist_ok=True)
    
    # Generate images for each spec
    for spec in image_specs:
        directory = spec['directory']
        filename = spec['filename']
        target_size = get_target_size(directory)
        
        if file_needs_generation(directory, filename):
            print(f"\nProcessing: {spec['description']}")
            success = generate_image(spec['prompt'], directory, filename)
            
            if success:
                print(f"✓ Successfully generated image for {spec['description']}")
                
                # Compress images to appropriate size
                full_path = os.path.join(directory, filename)
                compress_success = compress_image(full_path, target_size)
                if compress_success:
                    print(f"✓ Successfully compressed {filename} to {target_size[0]}x{target_size[1]}")
                else:
                    print(f"✗ Failed to compress {filename}")
            else:
                print(f"✗ Failed to generate image for {spec['description']}")
        else:
            print(f"\nSkipping {spec['description']} - file already exists and is not empty")
            
            # Check if existing image needs compression
            full_path = os.path.join(directory, filename)
            try:
                with Image.open(full_path) as img:
                    width, height = img.size
                    if width != target_size[0] or height != target_size[1]:
                        print(f"Resizing existing image {filename} to {target_size[0]}x{target_size[1]}...")
                        compress_success = compress_image(full_path, target_size)
                        if compress_success:
                            print(f"✓ Successfully resized existing {filename} to {target_size[0]}x{target_size[1]}")
                        else:
                            print(f"✗ Failed to resize existing {filename}")
            except Exception as e:
                print(f"Error checking image size for {filename}: {str(e)}")
    
    print("\nImage generation process completed!")    

if __name__ == "__main__":
    main()
